#ifndef __IANGELS_SAFETENSORS_LOADER_H__
#define __IANGELS_SAFETENSORS_LOADER_H__

#include <vector>
#include <string>

#include "status.h"
#include "config.h"

namespace iangels {

constexpr uint64_t MAX_HEADER_SIZE = 1024 * 1024 * 1024; // 1GB

class SafeTensorsLoader {

public:
    SafeTensorsLoader() = default;
    ~SafeTensorsLoader() = default;
    
    struct SafeTensor {
        std::string filename_;
        uintmax_t file_bytes_;

        Config header_;
        size_t header_bytes_;

        char* weights_addr_;
    };

    Ret load(const std::vector<std::string>& filenames);

    Ret load_header_size(std::ifstream& file, SafeTensor& tensor);
    Ret load_header(std::ifstream& file, SafeTensor& tensor);
    Ret mmap_weights(const std::string& file, SafeTensor& tensor);
    Ret unmap_weights();

    const Config& header(int index) const { return tensors_[index].header_; }
    const std::vector<SafeTensor>& tensors() const { return tensors_; }

private:
    std::vector<SafeTensor> tensors_{};
}; // class SafeTensorsLoader
    
} // namespace iangels

#endif // __IANGELS_SAFETENSORS_LOADER_H__